# 代码来源于以下的链接
# https://blog.csdn.net/Wei_sx/article/details/145055459
from keras.src import models, layers


def vgg_block(num_convs, num_filters):
    blk = models.Sequential()
    for _ in range(num_convs):
        blk.add(layers.Conv2D(num_filters, kernel_size=3, padding='same', activation='relu'))
    blk.add(layers.MaxPooling2D(pool_size=2, strides=2))
    return blk


def vgg(conv_arch):
    net = models.Sequential()
    # generate five blocks
    for (num_convs, num_filters) in conv_arch:
        net.add(vgg_block(num_convs, num_filters))

    net.add(models.Sequential([
        layers.Flatten(),
        layers.Dense(4096, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(4096, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(10, activation='softmax')
    ]))

    return net


conv_arch = ((2, 64), (2, 128), (3, 256), (3, 512), (3, 512))
model = vgg(conv_arch)

print(model)
print(444)
